import logging
from typing import Dict, Optional, List, Tuple
from utils.db_manager import DBManager
import aiohttp
import json

class TermTranslationService:
    """专业术语翻译服务"""
    
    def __init__(self):
        self.logger = logging.getLogger(self.__class__.__name__)
        self.db_manager = DBManager()
        self.api_keys = {
            'google': os.getenv('GOOGLE_TRANSLATE_API_KEY'),
            'deepl': os.getenv('DEEPL_API_KEY')
        }
        
    async def translate_term(self, term: str, source_lang: str, target_lang: str) -> Optional[str]:
        """翻译单个术语"""
        try:
            # 1. 首先查询本地术语库
            local_translation = await self._get_local_term(term, source_lang, target_lang)
            if local_translation:
                return local_translation
            
            # 2. 尝试使用专业翻译API
            api_translation = await self._translate_with_api(term, source_lang, target_lang)
            if api_translation:
                # 保存到本地术语库
                await self._save_term(term, api_translation, source_lang, target_lang)
                return api_translation
            
            # 3. 使用大模型翻译
            model_translation = await self._translate_with_model(term, source_lang, target_lang)
            if model_translation:
                # 保存到本地术语库
                await self._save_term(term, model_translation, source_lang, target_lang)
                return model_translation
            
            return None
            
        except Exception as e:
            self.logger.error(f"Term translation failed: {str(e)}")
            return None
    
    async def _get_local_term(self, term: str, source_lang: str, target_lang: str) -> Optional[str]:
        """从本地术语库查询"""
        try:
            conn = self.db_manager.get_connection()
            cursor = conn.cursor()
            
            query = """
                SELECT target_term, confidence 
                FROM technical_terms 
                WHERE source_term = %s 
                AND source_lang = %s 
                AND target_lang = %s 
                AND confidence >= 0.8  -- 只返回高可信度的术语
                ORDER BY confidence DESC, updated_at DESC
                LIMIT 1
            """
            
            cursor.execute(query, (term, source_lang, target_lang))
            result = cursor.fetchone()
            
            cursor.close()
            conn.close()
            
            return result[0] if result else None
            
        except Exception as e:
            self.logger.error(f"Local term query failed: {str(e)}")
            return None
    
    async def _translate_with_api(self, term: str, source_lang: str, target_lang: str) -> Optional[str]:
        """使用专业翻译API"""
        # 尝试使用Google Translate API
        try:
            async with aiohttp.ClientSession() as session:
                params = {
                    'q': term,
                    'source': source_lang,
                    'target': target_lang,
                    'key': self.api_keys['google']
                }
                async with session.get('https://translation.googleapis.com/language/translate/v2', params=params) as response:
                    if response.status == 200:
                        data = await response.json()
                        return data['data']['translations'][0]['translatedText']
        except Exception as e:
            self.logger.warning(f"Google Translate API failed: {str(e)}")
        
        # 尝试使用DeepL API
        try:
            async with aiohttp.ClientSession() as session:
                params = {
                    'text': term,
                    'source_lang': source_lang.upper(),
                    'target_lang': target_lang.upper(),
                    'auth_key': self.api_keys['deepl']
                }
                async with session.get('https://api-free.deepl.com/v2/translate', params=params) as response:
                    if response.status == 200:
                        data = await response.json()
                        return data['translations'][0]['text']
        except Exception as e:
            self.logger.warning(f"DeepL API failed: {str(e)}")
        
        return None
    
    async def _translate_with_model(self, term: str, source_lang: str, target_lang: str) -> Optional[str]:
        """使用本地模型翻译"""
        try:
            # 使用NLLB模型进行翻译
            from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
            
            model_name = "facebook/nllb-200-distilled-600M"
            tokenizer = AutoTokenizer.from_pretrained(model_name)
            model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
            
            # 转换语言代码
            source_lang = f"{source_lang}_Latn"
            target_lang = f"{target_lang}_Latn"
            
            # 生成翻译
            inputs = tokenizer(term, return_tensors="pt")
            translated_tokens = model.generate(
                **inputs,
                forced_bos_token_id=tokenizer.lang_code_to_id[target_lang],
                max_length=128,
                num_beams=5,
                length_penalty=1.0
            )
            
            translation = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
            return translation
            
        except Exception as e:
            self.logger.error(f"Model translation failed: {str(e)}")
            return None
    
    async def _save_term(self, source_term: str, target_term: str, 
                        source_lang: str, target_lang: str, 
                        confidence: float = 0.85) -> bool:
        """保存术语到数据库"""
        try:
            conn = self.db_manager.get_connection()
            cursor = conn.cursor()
            
            query = """
                INSERT INTO technical_terms 
                (source_term, target_term, source_lang, target_lang, confidence, source)
                VALUES (%s, %s, %s, %s, %s, %s)
                ON DUPLICATE KEY UPDATE
                confidence = GREATEST(confidence, VALUES(confidence)),
                updated_at = CURRENT_TIMESTAMP
            """
            
            cursor.execute(query, (
                source_term, target_term, source_lang, target_lang,
                confidence, 'api' if confidence > 0.9 else 'model'
            ))
            
            conn.commit()
            cursor.close()
            conn.close()
            
            return True
            
        except Exception as e:
            self.logger.error(f"Failed to save term: {str(e)}")
            return False 